//===--- FunctionSignatureOpts.cpp - Optimizes function signatures --------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
///
/// \file
///
/// This pass defines function signature related optimizations.
/// When a function signature optimization is performed, changes are made to
/// the original function and after all function signature optimizations are
/// finished, a new function is created and the old function is turned into
/// a thunk.
///
/// Another possibility is to implement these optimizations as separate passes,
/// but then we would send slightly different functions to the pass pipeline
/// multiple times through notifyPassManagerOfFunction.
///
/// TODO: Optimize function with generic parameters.
///
/// TODO: Improve epilogue release matcher, i.e. do a data flow instead of
/// only finding releases in the return block.
///
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "pil-function-signature-opt"

#include "polarphp/pil/optimizer/internal/funcsignaturetransforms/FunctionSignatureOpts.h"
#include "polarphp/pil/lang/DebugUtils.h"
#include "polarphp/pil/lang/PILCloner.h"
#include "polarphp/pil/lang/PILFunction.h"
#include "polarphp/pil/lang/PILValue.h"
#include "polarphp/pil/optimizer/analysis/CallerAnalysis.h"
#include "polarphp/pil/optimizer/analysis/EpilogueARCAnalysis.h"
#include "polarphp/pil/optimizer/analysis/RCIdentityAnalysis.h"
#include "polarphp/pil/optimizer/passmgr/Passes.h"
#include "polarphp/pil/optimizer/passmgr/Transforms.h"
#include "polarphp/pil/optimizer/utils/InstOptUtils.h"
#include "polarphp/pil/optimizer/utils/PILInliner.h"
#include "polarphp/pil/optimizer/utils/PILOptFunctionBuilder.h"
#include "polarphp/pil/optimizer/utils/SpecializationMangler.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"

using namespace polar;

STATISTIC(NumFunctionSignaturesOptimized, "Total func sig optimized");
STATISTIC(NumDeadArgsEliminated, "Total dead args eliminated");
STATISTIC(NumOwnedConvertedToGuaranteed, "Total owned args -> guaranteed args");
STATISTIC(NumOwnedConvertedToNotOwnedResult, "Total owned result -> not owned result");
STATISTIC(NumSROAArguments, "Total SROA arguments optimized");

using PILParameterInfoList = llvm::SmallVector<PILParameterInfo, 8>;
using ArgumentIndexMap = llvm::SmallDenseMap<int, int>;

//===----------------------------------------------------------------------===//
//                           Optimization Hueristic
//===----------------------------------------------------------------------===//

/// Set to true to enable the support for partial specialization.
static llvm::cl::opt<bool>
   FSOEnableGenerics("sil-fso-enable-generics", llvm::cl::init(true),
                     llvm::cl::desc("Support function signature optimization "
                                    "of generic functions"));

static llvm::cl::opt<bool>
   FSOOptimizeIfNotCalled("sil-fso-optimize-if-not-called",
                          llvm::cl::init(false),
                          llvm::cl::desc("Optimize even if a function isn't "
                                         "called. For testing only!"));

static bool isSpecializableRepresentation(PILFunctionTypeRepresentation Rep,
                                          bool OptForPartialApply) {
   switch (Rep) {
      case PILFunctionTypeRepresentation::Method:
      case PILFunctionTypeRepresentation::Closure:
      case PILFunctionTypeRepresentation::Thin:
      case PILFunctionTypeRepresentation::Thick:
      case PILFunctionTypeRepresentation::CFunctionPointer:
         return true;
      case PILFunctionTypeRepresentation::WitnessMethod:
         return OptForPartialApply;
      case PILFunctionTypeRepresentation::ObjCMethod:
      case PILFunctionTypeRepresentation::Block:
         return false;
   }

   llvm_unreachable("Unhandled PILFunctionTypeRepresentation in switch.");
}

/// Returns true if F is a function which the pass knows how to specialize
/// function signatures for.
static bool canSpecializeFunction(PILFunction *F,
                                  const CallerAnalysis::FunctionInfo *FuncInfo,
                                  bool OptForPartialApply) {
   // Do not specialize the signature of PILFunctions that are external
   // declarations since there is no body to optimize.
   if (F->isExternalDeclaration())
      return false;

   // For now ignore functions with indirect results.
   if (F->getConventions().hasIndirectPILResults())
      return false;

   // For now ignore coroutines.
   if (F->getLoweredFunctionType()->isCoroutine())
      return false;

   // Do not specialize the signature of always inline functions. We
   // will just inline them and specialize each one of the individual
   // functions that these sorts of functions are inlined into.
   // It is OK to specialize always inline functions if they are
   // used by partial_apply instructions.
   assert(!OptForPartialApply || FuncInfo);
   if (F->getInlineStrategy() == Inline_t::AlwaysInline &&
       (!OptForPartialApply || !FuncInfo->getMinPartialAppliedArgs()))
      return false;

   // For now ignore generic functions to keep things simple...
   if (!FSOEnableGenerics && F->getLoweredFunctionType()->isPolymorphic())
      return false;

   // Make sure F has a linkage that we can optimize.
   if (!isSpecializableRepresentation(F->getRepresentation(),
                                      OptForPartialApply))
      return false;

   return true;
}

//===----------------------------------------------------------------------===//
//                  Function Signature Transform Descriptor
//===----------------------------------------------------------------------===//

void FunctionSignatureTransformDescriptor::addThunkArgument(
   ArgumentDescriptor &AD, PILBuilder &Builder, PILBasicBlock *BB,
   llvm::SmallVectorImpl<PILValue> &NewArgs) {
   // Dead argument.
   if (AD.IsEntirelyDead) {
      return;
   }

   // Explode the argument.
   if (AD.Explode) {
      llvm::SmallVector<PILValue, 4> LeafValues;
      AD.ProjTree.createTreeFromValue(Builder, BB->getParent()->getLocation(),
                                      BB->getArgument(AD.Index), LeafValues);
      NewArgs.append(LeafValues.begin(), LeafValues.end());
      return;
   }

   // All other arguments get pushed as what they are.
   NewArgs.push_back(BB->getArgument(AD.Index));
}

std::string
FunctionSignatureTransformDescriptor::createOptimizedPILFunctionName() {
   PILFunction *F = OriginalFunction;

   auto P = demangling::SpecializationPass::FunctionSignatureOpts;
   mangle::FunctionSignatureSpecializationMangler Mangler(P, F->isSerialized(),
                                                          F);

   // Handle arguments' changes.
   for (unsigned i : indices(ArgumentDescList)) {
      const ArgumentDescriptor &Arg = ArgumentDescList[i];
      if (Arg.IsEntirelyDead) {
         Mangler.setArgumentDead(i);
         // No point setting other attribute if argument is dead.
         continue;
      }

      // If we have an @owned argument and found a callee release for it,
      // convert the argument to guaranteed.
      if (Arg.OwnedToGuaranteed) {
         Mangler.setArgumentOwnedToGuaranteed(i);
      }

      // If this argument is not dead and we can explode it, add 's' to the
      // mangling.
      if (Arg.Explode) {
         Mangler.setArgumentSROA(i);
      }
   }

   // Handle return value's change.
   // FIXME: handle multiple direct results here
   if (ResultDescList.size() == 1 && !ResultDescList[0].CalleeRetain.empty()) {
      Mangler.setReturnValueOwnedToUnowned();
   }

   return Mangler.mangle();
}

/// Collect all archetypes used by a function.
static bool usesGenerics(PILFunction *F,
                         ArrayRef<PILParameterInfo> InterfaceParams,
                         ArrayRef<PILResultInfo> InterfaceResults) {
   CanPILFunctionType FTy = F->getLoweredFunctionType();
   auto HasGenericSignature = FTy->getSubstGenericSignature() != nullptr;
   if (!HasGenericSignature)
      return false;

   bool UsesGenerics = false;

   auto FindArchetypesAndGenericTypes = [&UsesGenerics](Type Ty) {
      if (Ty.findIf([](Type Ty) -> bool {
         return (Ty->hasTypeParameter() || Ty->hasArchetype());
      }))
         UsesGenerics = true;
   };

   for (auto Param : InterfaceParams) {
      Param.getInterfaceType().visit(FindArchetypesAndGenericTypes);
   }

   for (auto Result : InterfaceResults) {
      Result.getInterfaceType().visit(FindArchetypesAndGenericTypes);
   }

   if (UsesGenerics)
      return UsesGenerics;

   for (auto &BB : *F) {
      for (auto &I : BB) {
         for (auto Arg : BB.getArguments()) {
            if (&BB != &*F->begin()) {
               // Scan types of all BB arguments. Ignore the entry BB, because
               // it is handled in a special way.
               Arg->getType().getAstType().visit(FindArchetypesAndGenericTypes);
               if (UsesGenerics)
                  return UsesGenerics;
            }
         }
         // Scan types of all operands.
         for (auto &Op : I.getAllOperands()) {
            Op.get()->getType().getAstType().visit(FindArchetypesAndGenericTypes);
         }
         // Scan all substitutions of apply instructions.
         if (auto AI = ApplySite::isa(&I)) {
            auto Subs = AI.getSubstitutionMap();
            for (auto Replacement : Subs.getReplacementTypes()) {
               Replacement.visit(FindArchetypesAndGenericTypes);
            }
         }
         // Scan all substitutions of builtin instructions.
         if (auto *BI = dyn_cast<BuiltinInst>(&I)) {
            auto Subs = BI->getSubstitutions();
            for (auto Ty : Subs.getReplacementTypes()) {
               Ty.visit(FindArchetypesAndGenericTypes);
            }
         }

         // Scan the result type of the instruction.
         for (auto V : I.getResults()) {
            V->getType().getAstType().visit(FindArchetypesAndGenericTypes);
         }

         if (UsesGenerics)
            return UsesGenerics;
      }
   }
   return UsesGenerics;
}

// Map the parameter, result and error types out of context to get the interface
// type.
static void mapInterfaceTypes(PILFunction *F,
                              MutableArrayRef<PILParameterInfo> InterfaceParams,
                              MutableArrayRef<PILResultInfo> InterfaceResults,
                              Optional<PILResultInfo> &InterfaceErrorResult) {

   for (auto &Param : InterfaceParams) {
      if (!Param.getInterfaceType()->hasArchetype())
         continue;
      Param = PILParameterInfo(
         Param.getInterfaceType()->mapTypeOutOfContext()->getCanonicalType(),
         Param.getConvention());
   }

   for (auto &Result : InterfaceResults) {
      if (!Result.getInterfaceType()->hasArchetype())
         continue;
      auto InterfaceResult = Result.getWithInterfaceType(
         Result.getInterfaceType()->mapTypeOutOfContext()->getCanonicalType());
      Result = InterfaceResult;
   }

   if (InterfaceErrorResult.hasValue()) {
      if (InterfaceErrorResult.getValue().getInterfaceType()->hasArchetype()) {
         InterfaceErrorResult =
            PILResultInfo(InterfaceErrorResult.getValue()
                             .getInterfaceType()
                             ->mapTypeOutOfContext()
                             ->getCanonicalType(),
                          InterfaceErrorResult.getValue().getConvention());
      }
   }
}

CanPILFunctionType
FunctionSignatureTransformDescriptor::createOptimizedPILFunctionType() {
   PILFunction *F = OriginalFunction;
   CanPILFunctionType FTy = F->getLoweredFunctionType();
   auto ExpectedFTy = F->getLoweredType().castTo<PILFunctionType>();
   auto HasGenericSignature = FTy->getSubstGenericSignature() != nullptr;

   // The only way that we modify the arity of function parameters is here for
   // dead arguments. Doing anything else is unsafe since by definition non-dead
   // arguments will have SSA uses in the function. We would need to be smarter
   // in our moving to handle such cases.
   llvm::SmallVector<PILParameterInfo, 8> InterfaceParams;
   for (auto &ArgDesc : ArgumentDescList) {
      computeOptimizedArgInterface(ArgDesc, InterfaceParams);
   }

   // ResultDescs only covers the direct results; we currently can't ever
   // change an indirect result.  Piece the modified direct result information
   // back into the all-results list.
   llvm::SmallVector<PILResultInfo, 8> InterfaceResults;
   for (PILResultInfo InterfaceResult : FTy->getResults()) {
      if (InterfaceResult.isFormalDirect()) {
         auto &RV = ResultDescList[0];
         if (!RV.CalleeRetain.empty()) {
            ++NumOwnedConvertedToNotOwnedResult;
            InterfaceResults.push_back(PILResultInfo(InterfaceResult.getInterfaceType(),
                                                     ResultConvention::Unowned));
            continue;
         }
      }

      InterfaceResults.push_back(InterfaceResult);
   }

   llvm::SmallVector<PILYieldInfo, 8> InterfaceYields;
   for (PILYieldInfo InterfaceYield : FTy->getYields()) {
      // For now, don't touch the yield types.
      InterfaceYields.push_back(InterfaceYield);
   }

   bool UsesGenerics = false;
   if (HasGenericSignature) {
      // Not all of the generic type parameters are used by the function
      // parameters.
      // Check which of the generic type parameters are not used and check if they
      // are used anywhere in the function body. If this is not the case, we can
      // remove the unused generic type parameters from the generic signature.
      // This makes the code both smaller and faster, because no implicit
      // parameters for type metadata and conformances need to be passed to the
      // callee at the LLVM IR level.
      // TODO: Implement a more precise analysis, so that we can eliminate only
      // those generic parameters which are not used.
      UsesGenerics = usesGenerics(F, InterfaceParams, InterfaceResults);

      // The set of used archetypes is complete now.
      if (!UsesGenerics) {
         // None of the generic type parameters are used.
         LLVM_DEBUG(llvm::dbgs() << "None of generic parameters are used by "
                                 << F->getName() << "\n";
                       llvm::dbgs() << "Interface params:\n";
                       for (auto Param : InterfaceParams) {
                          Param.getInterfaceType().dump(llvm::dbgs());
                       }

                       llvm::dbgs() << "Interface results:\n";
                       for (auto Result : InterfaceResults) {
                          Result.getInterfaceType().dump(llvm::dbgs());
                       });
      }
   }

   // Don't use a method representation if we modified self.
   auto ExtInfo = FTy->getExtInfo();
   auto witnessMethodConformance = FTy->getWitnessMethodConformanceOrInvalid();
   if (shouldModifySelfArgument) {
      ExtInfo = ExtInfo.withRepresentation(PILFunctionTypeRepresentation::Thin);
      witnessMethodConformance = InterfaceConformanceRef::forInvalid();
   }

   Optional<PILResultInfo> InterfaceErrorResult;
   if (ExpectedFTy->hasErrorResult()) {
      InterfaceErrorResult = ExpectedFTy->getErrorResult();
   }

   // Map the parameter, result and error types out of context to get the
   // proper interface type. This is required for generic functions.
   mapInterfaceTypes(F, InterfaceParams, InterfaceResults, InterfaceErrorResult);

   GenericSignature GenericSig =
      UsesGenerics ? FTy->getSubstGenericSignature() : nullptr;

   return PILFunctionType::get(
      GenericSig, ExtInfo, FTy->getCoroutineKind(), FTy->getCalleeConvention(),
      InterfaceParams, InterfaceYields, InterfaceResults, InterfaceErrorResult,
      FTy->getSubstitutions(), FTy->isGenericSignatureImplied(),
      F->getModule().getAstContext(), witnessMethodConformance);
}

/// Compute what the function interface will look like based on the
/// optimization we are doing on the given argument descriptor. Default
/// implementation simply passes it through.
void FunctionSignatureTransformDescriptor::computeOptimizedArgInterface(
   ArgumentDescriptor &AD, SmallVectorImpl<PILParameterInfo> &Out) {
   // If this argument is live, but we cannot optimize it.
   if (!AD.canOptimizeLiveArg()) {
      if (AD.PInfo.hasValue())
         Out.push_back(AD.PInfo.getValue());
      return;
   }

   // If we have a dead argument, bail.
   if (AD.IsEntirelyDead) {
      ++NumDeadArgsEliminated;
      return;
   }

   // Explode the argument or not ?
   if (AD.Explode) {
      ++NumSROAArguments;
      llvm::SmallVector<const ProjectionTreeNode *, 8> LeafNodes;
      AD.ProjTree.getLiveLeafNodes(LeafNodes);
      for (auto Node : LeafNodes) {
         PILType Ty = Node->getType();
         LLVM_DEBUG(llvm::dbgs() << "                " << Ty << "\n");
         // If Ty is trivial, just pass it directly.
         if (Ty.isTrivial(*AD.Arg->getFunction())) {
            PILParameterInfo NewInfo(Ty.getAstType(),
                                     ParameterConvention::Direct_Unowned);
            Out.push_back(NewInfo);
            continue;
         }

         // Ty is not trivial, pass it through as the original calling convention.
         auto ParameterConvention = AD.PInfo.getValue().getConvention();
         if (AD.OwnedToGuaranteed) {
            if (ParameterConvention == ParameterConvention::Direct_Owned)
               ParameterConvention = ParameterConvention::Direct_Guaranteed;
            else if (ParameterConvention == ParameterConvention::Indirect_In)
               ParameterConvention = ParameterConvention::Indirect_In_Guaranteed;
            else {
               llvm_unreachable("Unknown parameter convention transformation");
            }
         }
         PILParameterInfo NewInfo(Ty.getAstType(), ParameterConvention);
         Out.push_back(NewInfo);
      }
      return;
   }

   // If we cannot explode this value, handle callee release and return.
   // If we found releases in the callee in the last BB on an @owned
   // parameter, change the parameter to @guaranteed and continue...
   if (AD.OwnedToGuaranteed) {
      ++NumOwnedConvertedToGuaranteed;
      auto ParameterConvention = AD.PInfo.getValue().getConvention();
      if (ParameterConvention == ParameterConvention::Direct_Owned)
         ParameterConvention = ParameterConvention::Direct_Guaranteed;
      else if (ParameterConvention == ParameterConvention::Indirect_In)
         ParameterConvention = ParameterConvention::Indirect_In_Guaranteed;
      else {
         llvm_unreachable("Unknown parameter convention transformation");
      }

      PILParameterInfo NewInfo(AD.PInfo.getValue().getInterfaceType(),
                               ParameterConvention);
      Out.push_back(NewInfo);
      return;
   }

   // Otherwise just propagate through the parameter info.
   Out.push_back(AD.PInfo.getValue());
}

//===----------------------------------------------------------------------===//
//                        Function Signature Transform
//===----------------------------------------------------------------------===//

void FunctionSignatureTransform::createFunctionSignatureOptimizedFunction() {
   // Create the optimized function!
   PILFunction *F = TransformDescriptor.OriginalFunction;
   PILModule &M = F->getModule();
   std::string Name = TransformDescriptor.createOptimizedPILFunctionName();
   // The transformed function must not already exist. This would indicate
   // repeated application of FSO on the same function. That situation should be
   // detected earlier by avoiding reoptimization of FSO thunks.
   assert(!F->getModule().hasFunction(Name));

   PILLinkage linkage = getSpecializedLinkage(F, F->getLinkage());

   LLVM_DEBUG(llvm::dbgs() << "  -> create specialized function " << Name
                           << "\n");

   auto NewFTy = TransformDescriptor.createOptimizedPILFunctionType();
   GenericEnvironment *NewFGenericEnv;
   if (NewFTy->getSubstGenericSignature()) {
      NewFGenericEnv = F->getGenericEnvironment();
   } else {
      NewFGenericEnv = nullptr;
   }

   // The specialized function is an internal detail, so we need to disconnect it
   // from a parent class, if one exists, thus the override of the
   // classSubclassScope.
   TransformDescriptor.OptimizedFunction = FunctionBuilder.createFunction(
      linkage, Name, NewFTy, NewFGenericEnv, F->getLocation(), F->isBare(),
      F->isTransparent(), F->isSerialized(), IsNotDynamic, F->getEntryCount(),
      F->isThunk(),
      /*classSubclassScope=*/SubclassScope::NotApplicable,
      F->getInlineStrategy(), F->getEffectsKind(), nullptr, F->getDebugScope());
   PILFunction *NewF = TransformDescriptor.OptimizedFunction.get();
   if (!F->hasOwnership()) {
      NewF->setOwnershipEliminated();
   }

   if (F->isSpecialization()) {
      NewF->setSpecializationInfo(F->getSpecializationInfo());
   }

   // Then we transfer the body of F to NewF.
   NewF->spliceBody(F);

   // Array semantic clients rely on the signature being as in the original
   // version.
   for (auto &Attr : F->getSemanticsAttrs()) {
      if (!StringRef(Attr).startswith("array."))
         NewF->addSemanticsAttr(Attr);
   }

   // Do the last bit of work to the newly created optimized function.
   DeadArgumentFinalizeOptimizedFunction();
   ArgumentExplosionFinalizeOptimizedFunction();

   // Update the ownership kinds of function entry BB arguments.
   for (auto Arg : NewF->begin()->getPILFunctionArguments()) {
      PILType MappedTy = Arg->getType();
      auto Ownershipkind =
         ValueOwnershipKind(*NewF, MappedTy, Arg->getArgumentConvention());
      Arg->setOwnershipKind(Ownershipkind);
   }

   // Create the thunk body !
   F->setThunk(IsSignatureOptimizedThunk);
   // The thunk now carries the information on how the signature is
   // optimized. If we inline the thunk, we will get the benefit of calling
   // the signature optimized function without additional setup on the
   // caller side.
   F->setInlineStrategy(AlwaysInline);
   PILBasicBlock *ThunkBody = F->createBasicBlock();
   for (auto &ArgDesc : TransformDescriptor.ArgumentDescList) {
      ThunkBody->createFunctionArgument(ArgDesc.Arg->getType(), ArgDesc.Decl);
   }

   PILLocation Loc = RegularLocation::getAutoGeneratedLocation();
   PILBuilder Builder(ThunkBody);
   Builder.setCurrentDebugScope(ThunkBody->getParent()->getDebugScope());

   FunctionRefInst *FRI = Builder.createFunctionRef(Loc, NewF);

   // Create the args for the thunk's apply, ignoring any dead arguments.
   llvm::SmallVector<PILValue, 8> ThunkArgs;
   for (auto &ArgDesc : TransformDescriptor.ArgumentDescList) {
      TransformDescriptor.addThunkArgument(ArgDesc, Builder, ThunkBody,
                                           ThunkArgs);
   }

   PILValue ReturnValue;
   PILType LoweredType = NewF->getLoweredType();
   PILType ResultType = NewF->getConventions().getPILResultType();
   auto GenCalleeType = NewF->getLoweredFunctionType();
   auto SubstCalleePILType = LoweredType;
   SubstitutionMap Subs;
   // Handle generic functions.
   if (GenCalleeType->isPolymorphic()) {
      // Produce a substitutions list and a set of substituted PIL types
      // required for creating a new PIL function.
      Subs = F->getForwardingSubstitutionMap();
      auto SubstCalleeType = GenCalleeType->substGenericArgs(
         M, Subs, Builder.getTypeExpansionContext());
      SubstCalleePILType = PILType::getPrimitiveObjectType(SubstCalleeType);
      PILFunctionConventions Conv(SubstCalleeType, M);
      ResultType = Conv.getPILResultType();
   }
   auto FunctionTy = LoweredType.castTo<PILFunctionType>();
   if (FunctionTy->hasErrorResult()) {
      // We need a try_apply to call a function with an error result.
      PILFunction *Thunk = ThunkBody->getParent();
      PILBasicBlock *NormalBlock = Thunk->createBasicBlock();
      ReturnValue =
         NormalBlock->createPhiArgument(ResultType, ValueOwnershipKind::Owned);
      PILBasicBlock *ErrorBlock = Thunk->createBasicBlock();
      PILType Error =
         PILType::getPrimitiveObjectType(FunctionTy->getErrorResult().getInterfaceType());
      auto *ErrorArg =
         ErrorBlock->createPhiArgument(Error, ValueOwnershipKind::Owned);
      Builder.createTryApply(Loc, FRI, Subs, ThunkArgs, NormalBlock, ErrorBlock);

      Builder.setInsertionPoint(ErrorBlock);
      Builder.createThrow(Loc, ErrorArg);
      Builder.setInsertionPoint(NormalBlock);
   } else {
      ReturnValue = Builder.createApply(Loc, FRI, Subs, ThunkArgs);
   }

   // Set up the return results.
   if (NewF->isNoReturnFunction()) {
      Builder.createUnreachable(Loc);
   } else {
      Builder.createReturn(Loc, ReturnValue);
   }

   // Do the last bit work to finalize the thunk.
   OwnedToGuaranteedFinalizeThunkFunction(Builder, F);

   assert(F->getDebugScope()->Parent != NewF->getDebugScope()->Parent);
}

// Run the optimization.
bool FunctionSignatureTransform::run(bool hasCaller) {
   // We use a reference here on purpose so our transformations can know if we
   // are going to make a thunk and thus should just optimize.
   bool &Changed = TransformDescriptor.Changed;
   bool hasOnlyDirectInModuleCallers =
      TransformDescriptor.hasOnlyDirectInModuleCallers;
   PILFunction *F = TransformDescriptor.OriginalFunction;

   // Never repeat the same function signature optimization on the same function.
   // Multiple function signature optimizations are composed by successively
   // optmizing the newly created functions. Each optimization creates a new
   // level of thunk. Those should all be ultimately inlined away.
   //
   // This happens, for example, when a new reference to the original function is
   // discovered during devirtualization. That will cause the original function
   // (now and FSO thunk) to be pushed back on the function pass pipeline.
   if (F->isThunk() == IsSignatureOptimizedThunk) {
      LLVM_DEBUG(llvm::dbgs() << "  FSO already performed on this thunk\n");
      return false;
   }

   // If we are asked to assume a caller for testing purposes, set the flag.
   hasCaller |= FSOOptimizeIfNotCalled;

   if (!hasCaller && (F->getDynamicallyReplacedFunction() ||
                      canBeCalledIndirectly(F->getRepresentation()))) {
      LLVM_DEBUG(llvm::dbgs() << "  function has no caller -> abort\n");
      return false;
   }

   // Run OwnedToGuaranteed optimization.
   if (OwnedToGuaranteedAnalyze()) {
      Changed = true;
      LLVM_DEBUG(llvm::dbgs() << "  transform owned-to-guaranteed\n");
      OwnedToGuaranteedTransform();
   }

   // Run DeadArgument elimination transformation. We only specialize
   // if this function has a caller inside the current module or we have
   // already created a thunk.
   if ((hasCaller || Changed || hasOnlyDirectInModuleCallers) &&
       DeadArgumentAnalyzeParameters()) {
      Changed = true;
      LLVM_DEBUG(llvm::dbgs() << "  remove dead arguments\n");
      DeadArgumentTransformFunction();
   }

   // Run ArgumentExplosion transformation. We only specialize
   // if this function has a caller inside the current module or we have
   // already created a thunk.
   //
   // NOTE: we run argument explosion last because we've already initialized
   // the ArgumentDescList to have unexploded number of arguments. Exploding
   // it without changing the argument count is not going to help with
   // owned-to-guaranteed transformation.
   //
   // In order to not miss any opportunity, we send the optimized function
   // to the passmanager to optimize any opportunities exposed by argument
   // explosion.
   if ((hasCaller || Changed || hasOnlyDirectInModuleCallers) &&
       ArgumentExplosionAnalyzeParameters()) {
      Changed = true;
   }

   // Check if generic signature of the function could be changed by
   // removed some unused generic arguments.
   if (F->getLoweredFunctionType()->isPolymorphic() &&
       TransformDescriptor.createOptimizedPILFunctionType() !=
       F->getLoweredFunctionType()) {
      Changed = true;
   }

   // Create the specialized function and invalidate the old function.
   if (Changed) {
      createFunctionSignatureOptimizedFunction();
   }
   return Changed;
}

// Run dead argument elimination of partially applied functions.
//
// After this optimization CapturePropagation can replace the partial_apply by a
// direct reference to the specialized function.
bool FunctionSignatureTransform::removeDeadArgs(int minPartialAppliedArgs) {
   if (minPartialAppliedArgs < 1)
      return false;

   if (!DeadArgumentAnalyzeParameters())
      return false;

   PILFunction *F = TransformDescriptor.OriginalFunction;
   auto ArgumentDescList = TransformDescriptor.ArgumentDescList;

   // Check if at least the minimum number of partially applied arguments
   // are dead. Otherwise no partial_apply can be removed anyway.
   unsigned Size = ArgumentDescList.size();
   for (unsigned Idx : range(Size)) {
      if (Idx < Size - minPartialAppliedArgs) {
         // Don't remove arguments other than the partial applied ones, even if
         // they are dead.
         ArgumentDescList[Idx].IsEntirelyDead = false;
         continue;
      }

      // Is the partially applied argument dead?
      if (!ArgumentDescList[Idx].IsEntirelyDead)
         return false;

      // Currently we require that all dead parameters have trivial types.  The
      // reason is that it's very hard to find places where we can release those
      // parameters (as a replacement for the removed partial_apply).
      //
      // TODO: Maybe we can skip this restriction when we have semantic ARC.
      if (ArgumentDescList[Idx].Arg->getType().isTrivial(*F))
         continue;
      return false;
   }

   LLVM_DEBUG(llvm::dbgs() << "  remove dead arguments for partial_apply\n");
   DeadArgumentTransformFunction();
   createFunctionSignatureOptimizedFunction();
   return true;
}

//===----------------------------------------------------------------------===//
//                           Top Level Entry Point
//===----------------------------------------------------------------------===//

namespace {

class FunctionSignatureOpts : public PILFunctionTransform {

   /// If true, perform a special kind of dead argument elimination to enable
   /// removal of partial_apply instructions where all partially applied
   /// arguments are dead.
   bool OptForPartialApply;

public:

   FunctionSignatureOpts(bool OptForPartialApply) :
      OptForPartialApply(OptForPartialApply) { }

   void run() override {
      auto *F = getFunction();

      // Don't run function signature optimizations at -Os.
      if (F->optimizeForSize())
         return;

      // Don't optimize callees that should not be optimized.
      if (!F->shouldOptimize())
         return;

      if (F->isDynamicallyReplaceable())
         return;

      // This is the function to optimize.
      LLVM_DEBUG(llvm::dbgs() << "*** FSO on function: " << F->getName()
                              << " ***\n");

      // Check the signature of F to make sure that it is a function that we
      // can specialize. These are conditions independent of the call graph.
      // No need for CallerAnalysis if we are not optimizing for partial
      // applies.
      if (!OptForPartialApply &&
          !canSpecializeFunction(F, nullptr, OptForPartialApply)) {
         LLVM_DEBUG(llvm::dbgs() << "  cannot specialize function -> abort\n");
         return;
      }

      const CallerAnalysis *CA = PM->getAnalysis<CallerAnalysis>();
      const CallerAnalysis::FunctionInfo &FuncInfo = CA->getFunctionInfo(F);

      // Check the signature of F to make sure that it is a function that we
      // can specialize. These are conditions independent of the call graph.
      if (OptForPartialApply &&
          !canSpecializeFunction(F, &FuncInfo, OptForPartialApply)) {
         LLVM_DEBUG(llvm::dbgs() << "  cannot specialize function -> abort\n");
         return;
      }

      // Ok, we think we can perform optimization. Now perform a quick check
      auto *RCIA = getAnalysis<RCIdentityAnalysis>();
      auto *EA = PM->getAnalysis<EpilogueARCAnalysis>();

      // As we optimize the function more and more, the name of the function is
      // going to change, make sure the mangler is aware of all the changes done
      // to the function.
      auto P = demangling::SpecializationPass::FunctionSignatureOpts;
      mangle::FunctionSignatureSpecializationMangler Mangler(P,
                                                             F->isSerialized(), F);

      /// Keep a map between the exploded argument index and the original argument
      /// index.
      llvm::SmallDenseMap<int, int> AIM;
      int asize = F->begin()->getArguments().size();
      for (unsigned i : range(asize)) {
         AIM[i] = i;
      }

      // Allocate the argument and result descriptors.
      llvm::SpecificBumpPtrAllocator<ProjectionTreeNode> Allocator;
      llvm::SmallVector<ArgumentDescriptor, 4> ArgumentDescList;
      llvm::SmallVector<ResultDescriptor, 4> ResultDescList;
      auto Args = F->begin()->getPILFunctionArguments();
      for (unsigned i : indices(Args)) {
         ArgumentDescList.emplace_back(Args[i], Allocator);
      }
      for (PILResultInfo IR : F->getLoweredFunctionType()->getResults()) {
         ResultDescList.emplace_back(IR);
      }

      PILOptFunctionBuilder FuncBuilder(*this);
      // Owned to guaranteed optimization.
      FunctionSignatureTransform FST(FuncBuilder, F, RCIA, EA, Mangler, AIM,
                                     ArgumentDescList, ResultDescList,
                                     FuncInfo.foundAllCallers());

      bool Changed = false;
      if (OptForPartialApply) {
         Changed = FST.removeDeadArgs(FuncInfo.getMinPartialAppliedArgs());
      } else {
         Changed = FST.run(FuncInfo.hasDirectCaller());
      }

      if (!Changed) {
         return;
      }

      ++NumFunctionSignaturesOptimized;
      // The old function must be a thunk now.
      assert(F->isThunk() && "Old function should have been turned into a thunk");

      invalidateAnalysis(PILAnalysis::InvalidationKind::Everything);

      // Make sure the PM knows about this function. This will also help us
      // with self-recursion.
      addFunctionToPassManagerWorklist(FST.getOptimizedFunction(), F);

      if (!OptForPartialApply) {
         // We have to restart the pipeline for this thunk in order to run the
         // inliner (and other opts) again. This is important if the new
         // specialized function (which is called from this thunk) is
         // function-signature-optimized again and also becomes an
         // always-inline-thunk.
         restartPassPipeline();
      }
   }

};

} // end anonymous namespace

PILTransform *polar::createFunctionSignatureOpts() {
   return new FunctionSignatureOpts(/* OptForPartialApply */ false);
}

PILTransform *polar::createDeadArgSignatureOpt() {
   return new FunctionSignatureOpts(/* OptForPartialApply */ true);
}
